import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error
import matplotlib.pyplot as plt
import joblib

# 线性回归模型学习

# 数据准备
# 4×2矩阵，4个样本，2个特征
X = np.array([[1, 1], [1, 2], [2, 2], [2, 3], [3, 4],[4, 5], [5, 6], [6, 7]])
# 2×1矩阵, 矩阵乘法的变换系数
# y 标签
y = np.dot(X, np.array([1, 2])) + 3

# 数据分割
# test_size
#
# 指定测试集的比例或样本数量。可以是浮点数、整数或 None。
# 如果是浮点数（如 0.2），表示测试集的比例，即 20% 的数据会被分配到测试集中。
# 如果是整数（如 10），表示测试集的样本数量，即测试集中会有 10 个样本。
# 如果未指定，默认值为 0.25，即 25% 的数据会被分配到测试集中。
# random_state
#
# 控制随机数生成以确保结果的可重复性。可以是整数或 None。
# 如果指定了整数（如 42），每次运行代码都会得到相同的训练集和测试集分割结果，确保实验的可重复性。
# 如果未指定或为 None，每次分割的数据集都会不同。
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 模型选择
model = LinearRegression()

# 模型训练
model.fit(X_train, y_train)

# 模型评估
y_pred = model.predict(X_test)
mse = mean_squared_error(y_test, y_pred)
print(f"均方误差: {mse}")

# 画出训练数据和拟合的平面
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(X_train[:, 0], X_train[:, 1], y_train, color='blue', label='Training data')
ax.scatter(X_test[:, 0], X_test[:, 1], y_test, color='green', label='Testing data')

# 生成平面的网格点
x1_range = np.linspace(X_train[:, 0].min(), X_train[:, 0].max(), 10)
x2_range = np.linspace(X_train[:, 1].min(), X_train[:, 1].max(), 10)
x1_grid, x2_grid = np.meshgrid(x1_range, x2_range)
y_grid = model.intercept_ + model.coef_[0] * x1_grid + model.coef_[1] * x2_grid

# 绘制平面
ax.plot_surface(x1_grid, x2_grid, y_grid, alpha=0.5, color='red')

ax.set_xlabel('Feature 1')
ax.set_ylabel('Feature 2')
ax.set_zlabel('Target')
ax.set_title('Linear Regression')

plt.legend()
plt.show()

# 保存模型
joblib_file = "linear_regression_model.pkl"
joblib.dump(model, joblib_file)
print(f"模型已保存到 {joblib_file}")

# 加载模型
loaded_model = joblib.load(joblib_file)

# 进行预测
new_data = np.array([[3, 5]])
prediction = loaded_model.predict(new_data)
print(f"预测值: {prediction}")
